#!/usr/bin/env python3

import cv2 as cv
import numpy as np
from matplotlib import pyplot as plt

img = cv.imread('../pic/7.png', 0)

# 使用 np.fft.fft2 计算二维傅里叶变换
f = np.fft.fft2(img)
# 将傅里叶变换的结果移到中心
fshift = np.fft.fftshift(f)
# 计算幅度谱 计算傅里叶变换结果的幅度，并使用对数尺度来增强可视化效果
magnitude_spectrum = 20 * np.log(np.abs(fshift))

plt.subplot(121), plt.imshow(img, cmap='gray')
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
plt.subplot(122), plt.imshow(magnitude_spectrum, cmap='gray')
plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
plt.show(block=False)
# 等待键盘输入后关闭窗口
plt.waitforbuttonpress()
plt.close()

# 高通滤波器  逆傅里叶变换
rows, cols = img.shape
# 整除 获得中心坐标
crow, ccol = rows // 2, cols // 2
# 取消中心区域的低频部分
fshift[crow - 30:crow + 31, ccol - 30:ccol + 31] = 0
# 反向移动，DC分量回到左上角
f_ishift = np.fft.ifftshift(fshift)
# 逆傅里叶变换，得到滤除低频后的图像
img_back = np.fft.ifft2(f_ishift)
# 傅里叶变换中可能会存在复数，需要提取实部
img_back = np.real(img_back)

plt.subplot(121), plt.imshow(img, cmap='gray')
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
plt.subplot(122), plt.imshow(img_back, cmap='gray')
plt.title('Image after HPF'), plt.xticks([]), plt.yticks([])

plt.show(block=False)
plt.waitforbuttonpress()
plt.close()
